package org.basis.algorithm.tree.practice;

/**
 * 前缀树练习
 *
 * @author Mr_wenpan@163.com 2021/12/27 15:37
 */
public class Trie2 {

    TrieNode root;

    public Trie2() {
        root = new TrieNode();
    }

    /**
     * 插入一个数到前缀树上
     */
    public void insert(String word) {
        if (word == null) {
            return;
        }
        char[] chars = word.toCharArray();
        TrieNode cur = root;
        cur.pass++;

        for (int i = 0; i < chars.length; i++) {
            int index = chars[i] - 'a';
            // 没有走向该字符的路，则新建
            if (cur.nexts[index] == null) {
                cur.nexts[index] = new TrieNode();
            }
            cur.nexts[index].pass++;
            cur = cur.nexts[index];
        }
        cur.end++;
    }

    /**
     * 从前缀树上删除一个字符串
     */
    public void delete(String word) {
        if (word == null) {
            return;
        }
        // 先查询一下是否存在于前缀树上
        if (search(word) <= 0) {
            return;
        }
        TrieNode cur = root;
        char[] chars = word.toCharArray();
        for (int i = 0; i < chars.length; i++) {
            int index = chars[i] - 'a';
            // 把该字符的节点的pass--，如果--后等于0，则释放后面的节点（help gc）
            if (--cur.nexts[index].pass == 0) {
                cur.nexts[index] = null;
                return;
            }
            cur = cur.nexts[index];
        }
        // 以该字符结尾的字符串个数--
        cur.end--;
    }

    /**
     * 前缀树上是否存在以Word为前缀的
     * 返回前缀树上有多少个字符串是以Word为前缀的
     */
    public int prefix(String word) {
        if (word == null) {
            return 0;
        }
        char[] chars = word.toCharArray();
        TrieNode cur = root;
        for (int i = 0; i < chars.length; i++) {
            int index = chars[i] - 'a';
            if (cur.nexts[index] == null) {
                return -1;
            }
            cur = cur.nexts[index];
        }
        return cur.pass;
    }

    /**
     * 搜索一个字符串是否存在于前缀树上
     */
    public int search(String word) {
        if (word == null) {
            return -1;
        }
        TrieNode cur = root;
        char[] chars = word.toCharArray();
        for (int i = 0; i < chars.length; i++) {
            int index = chars[i] - 'a';
            if (cur.nexts[index] == null) {
                return -1;
            }
            cur = cur.nexts[index];
        }
        return cur.end;
    }

    /**
     * 前缀树节点
     */
    class TrieNode {
        public int pass;
        public int end;
        public TrieNode[] nexts;
        // 数据量大的情况下这里可以用hash表
//        public HashSet<String> nexts;

        public TrieNode() {
            pass = 0;
            end = 0;
            // 由于只有26个字符（只算小写字母），所以使用数组就够了，如果数据量非常大，那么这里改为hash表
            nexts = new TrieNode[26];
        }
    }

}


